import json
import os
import sys
from datetime import datetime
from typing import Dict, Any

from tqsdk import TqApi, TqAuth, TqKq
from rich.console import Console
from rich.table import Table
from rich.theme import Theme

CONFIG_PATH = r"f:\\Private.Net\\WorkShop\\Project_greewarm\\FastApi_tqsdk\\db_config.json"

# Rich 控制台主题
custom_theme = Theme({
    "title": "bold bright_cyan",
    "info": "bright_cyan",
    "warning": "bright_yellow",
    "error": "bright_red",
    "success": "bright_green",
    "dim": "grey50"
})
console = Console(theme=custom_theme)


def load_credentials():
    """尝试从配置文件读取账号信息，不存在则回退到默认测试账号。"""
    user = "cps168"
    password = "alibaba"
    if os.path.isfile(CONFIG_PATH):
        try:
            with open(CONFIG_PATH, "r", encoding="utf-8") as f:
                cfg = json.load(f)
            user = str(cfg.get("tq_user", user))
            password = str(cfg.get("tq_password", password))
            console.print(f"从配置文件 {CONFIG_PATH} 加载账号信息成功", style="info")
        except Exception as e:
            console.print(f"读取配置文件失败，使用默认账号: {e}", style="warning")
    else:
        console.print(f"未发现配置文件 {CONFIG_PATH}，使用默认账号", style="warning")
    return user, password


def format_money(v):
    try:
        return f"{float(v):,.2f}"
    except Exception:
        return str(v)


def collect_positions_from_obj(symbol: str, p: Any) -> Dict[str, Any]:
    """将持仓对象或原始字典统一为结构化字典"""
    if p is None:
        p = {}
    return {
        "symbol": symbol,
        "pos_long": (getattr(p, "pos_long", None) if hasattr(p, "pos_long") else p.get("pos_long", 0)) or 0,
        "pos_short": (getattr(p, "pos_short", None) if hasattr(p, "pos_short") else p.get("pos_short", 0)) or 0,
        "pos_long_his": (getattr(p, "pos_long_his", None) if hasattr(p, "pos_long_his") else p.get("pos_long_his", 0)) or 0,
        "pos_short_his": (getattr(p, "pos_short_his", None) if hasattr(p, "pos_short_his") else p.get("pos_short_his", 0)) or 0,
        "pos_long_today": (getattr(p, "pos_long_today", None) if hasattr(p, "pos_long_today") else p.get("pos_long_today", 0)) or 0,
        "pos_short_today": (getattr(p, "pos_short_today", None) if hasattr(p, "pos_short_today") else p.get("pos_short_today", 0)) or 0,
    }


def get_positions_map(api: TqApi) -> Dict[str, Dict[str, Any]]:
    """
    获取持仓的综合视图：优先 api.get_position()，并融合 api._data 的 accounts/trade/positions 兜底。
    不主动等待同步，仅使用当前可用缓存。
    """
    merged: Dict[str, Dict[str, Any]] = {}

    # 1) 直接使用 api.get_position() 无参字典
    pos_map = None
    try:
        pos_map = api.get_position()
    except Exception:
        pos_map = None
    if isinstance(pos_map, dict) and pos_map:
        for symbol, p in pos_map.items():
            merged[symbol] = collect_positions_from_obj(symbol, p)

    # 2) 从 api._data 收集候选
    accounts = api._data.get('accounts', {}) if hasattr(api, '_data') else {}
    trade_root = api._data.get('trade', {}) if hasattr(api, '_data') else {}
    top_positions = api._data.get('positions', {}) if hasattr(api, '_data') else {}

    # 从 accounts 路径
    if isinstance(accounts, dict) and accounts:
        # 选择一个账户键作为参考
        acct_key = None
        try:
            keys = list(accounts.keys())
            if keys:
                acct_key = keys[0]
        except Exception:
            acct_key = None
        if acct_key:
            acct_pos = accounts.get(acct_key, {}).get('positions', {})
            if isinstance(acct_pos, dict):
                for s, raw in acct_pos.items():
                    if s not in merged:
                        merged[s] = collect_positions_from_obj(s, raw)

    # 从 trade 路径
    if isinstance(trade_root, dict) and trade_root:
        for acct_key, node in trade_root.items():
            pos_map2 = node.get('positions', {}) if isinstance(node, dict) else {}
            if isinstance(pos_map2, dict):
                for s, raw in pos_map2.items():
                    if s not in merged:
                        merged[s] = collect_positions_from_obj(s, raw)

    # 从顶层 positions 路径
    if isinstance(top_positions, dict) and top_positions:
        for s, raw in top_positions.items():
            if s not in merged:
                merged[s] = collect_positions_from_obj(s, raw)

    return merged


def print_account_info(account, user: str):
    table = Table(title="账户资金信息", title_style="title")
    table.add_column("项目", style="dim")
    table.add_column("数值", justify="right")

    table.add_row("账户", str(user))
    table.add_row("更新时间", datetime.now().strftime("%Y-%m-%d %H:%M:%S"))
    table.add_row("可用资金", format_money(getattr(account, "available", 0)))
    table.add_row("资金权益", format_money(getattr(account, "balance", 0)))
    table.add_row("动态权益", format_money(getattr(account, "equity", 0)))
    table.add_row("当日盈亏", format_money(getattr(account, "close_profit", 0)))
    table.add_row("浮动盈亏", format_money(getattr(account, "float_profit", 0)))
    table.add_row("风险度", f"{getattr(account, 'risk_ratio', 0)}")
    table.add_row("保证金占用", format_money(getattr(account, "margin", 0)))
    table.add_row("手续费", format_money(getattr(account, "commission", 0)))

    console.print(table)


def print_positions(merged: Dict[str, Dict[str, Any]]):
    table = Table(title="账户持仓信息", title_style="title")
    table.add_column("品种代码", style="dim")
    table.add_column("总持仓", justify="right")
    table.add_column("多仓", justify="right")
    table.add_column("空仓", justify="right")
    table.add_column("昨仓", justify="right")
    table.add_column("今仓", justify="right")

    non_zero = 0
    # 排序：按总持仓降序
    symbols = sorted(merged.keys(), key=lambda s: (merged[s]["pos_long"] + merged[s]["pos_short"]), reverse=True)
    for symbol in symbols:
        p = merged[symbol]
        total_pos = (p["pos_long"] or 0) + (p["pos_short"] or 0)
        total_his = (p["pos_long_his"] or 0) + (p["pos_short_his"] or 0)
        total_today = (p["pos_long_today"] or 0) + (p["pos_short_today"] or 0)
        table.add_row(symbol, str(total_pos), str(p["pos_long"]), str(p["pos_short"]), str(total_his), str(total_today))
        if total_pos > 0 or total_today > 0:
            non_zero += 1

    console.print(table)
    if non_zero > 0:
        console.print(f"直接获取到 {non_zero} 个非零持仓", style="success")
    else:
        console.print("当前账户没有持仓或暂未返回持仓字典。", style="warning")


def main():
    console.print("天勤账户资金和持仓信息显示工具", style="title")
    console.print("=======================", style="dim")
    user, password = load_credentials()
    console.print(f"正在登录天勤账号: {user}", style="info")

    api = None
    try:
        # 立即连接，不等待同步，直接查询
        api = TqApi(account=TqKq(), auth=TqAuth(user, password))

        # 直接获取账户与持仓（不等待）
        account = api.get_account()
        merged_positions = get_positions_map(api)

        # 打印结果
        print_account_info(account, user)
        print_positions(merged_positions)

    except Exception as e:
        console.print(f"查询过程中出现错误: {e}", style="error")
        import traceback
        traceback.print_exc()
    finally:
        try:
            if api:
                api.close()
        except Exception:
            pass


if __name__ == "__main__":
    # 允许通过命令行参数覆盖默认配置文件路径（可选）
    if len(sys.argv) > 1:
        CONFIG_PATH = sys.argv[1]
    main()
